package dd.lo.flann;

import dd.lo.HelloCV;
import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.scene.text.Text;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.opencv.calib3d.Calib3d;
import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.DMatch;
import org.opencv.core.KeyPoint;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.core.MatOfDMatch;
import org.opencv.core.MatOfKeyPoint;
import org.opencv.core.MatOfPoint;
import org.opencv.core.MatOfPoint2f;
import org.opencv.core.Point;
import org.opencv.core.Rect;
import org.opencv.core.Scalar;
import org.opencv.core.Size;
import org.opencv.features2d.BFMatcher;
import org.opencv.features2d.Features2d;
import org.opencv.features2d.FlannBasedMatcher;
import org.opencv.features2d.ORB;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

public class BfMatchExample extends Application {
    public static void main(String[] args) {
        System.load(HelloCV.class.getResource("/libopencv_java481.dylib").getPath());
        BfMatchExample.start(args);
    }
    public static void start(String[] args) {
        launch(args);
    }

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private ImageView imageView;

    private String imgPath1;
    private String imgPath2;

    private Mat img1;
    private Mat img2;
    private Text result;

    private void compareImages() {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                img1 = Imgcodecs.imread(imgPath1, Imgcodecs.IMREAD_UNCHANGED);
                img2 = Imgcodecs.imread(imgPath2, Imgcodecs.IMREAD_UNCHANGED);
                adaptImg(img1);
                adaptImg(img2);
                //Resize img1 to fit the size of img2
//                ImageUtils.fitImgSize(img1, img2.width(), img2.height());
                Mat gray1 = new Mat(), gray2 = new Mat();
                Imgproc.cvtColor(img1, gray1, Imgproc.COLOR_BGR2GRAY);
                Imgproc.cvtColor(img2, gray2, Imgproc.COLOR_BGR2GRAY);
//                int kernelSize = 5;
//                Imgproc.blur(gray1, gray1, new Size(kernelSize, kernelSize)); //使用一个内核进行降噪
//                Imgproc.blur(gray2, gray2, new Size(kernelSize, kernelSize)); //使用一个内核进行降噪
//                //锐化图像
//                Mat kernel = new Mat(3, 3, CvType.CV_32SC1);
////                kernel.put(0, 0, new int[]{0, -1,  0, -1,  5, -1, 0, -1,  0});
//                kernel.put(0, 0, new int[]{-1, -1,  -1, -1,  9, -1, -1, -1,  -1});
//                Imgproc.filter2D(gray1, gray1, gray1.depth(), kernel);
//                Imgproc.filter2D(gray2, gray2, gray2.depth(), kernel);
//                SURF detector = SURF.create(); //Patented, set OPENCV_ENABLE_NONFREE CMake option and rebuild the library in function 'create'
                ORB detector = ORB.create();
                MatOfKeyPoint kp1 = new MatOfKeyPoint();
                //Descriptor must be in type of CV_32F. Ref: https://answers.opencv.org/question/11209/unsupported-format-or-combination-of-formats-in-buildindex-using-flann-algorithm/
                Mat des1 = new Mat();
                detector.detectAndCompute(gray1, new Mat(), kp1, des1);
                des1.convertTo(des1, CvType.CV_32F);
                MatOfKeyPoint kp2 = new MatOfKeyPoint();
                Mat des2 = new Mat();
                detector.detectAndCompute(gray2, new Mat(), kp2, des2);
                des2.convertTo(des2, CvType.CV_32F);
//                BFMatcher matcher = BFMatcher.create();
                BFMatcher matcher = BFMatcher.create(BFMatcher.BRUTEFORCE_HAMMING, true);
//                matcher.add(Collections.singletonList(des2));
//                List<MatOfDMatch> matches = new ArrayList<>();
//                matcher.knnMatch(des1, matches, 2);
//                //Filter good matches
//                List<MatOfDMatch> goodMatches = new ArrayList<>(matches.size());
//                for (MatOfDMatch match : matches) {
//                    DMatch[] dMatches = match.toArray();
//                    DMatch p1 = dMatches[0], p2 = dMatches[1];
//                    if (p1.distance < 0.7 * p2.distance) {
//                        goodMatches.add(match);
//                    }
//                }
//                Mat dst = new Mat();
//                Features2d.drawMatchesKnn(img1, kp1, img2, kp2, goodMatches, dst);
                MatOfDMatch matches = new MatOfDMatch();
                matcher.match(des1, des2, matches);
                List<DMatch> matchList = matches.toList();
                matchList.sort(Comparator.comparingDouble(m -> m.distance));
                //下面对所有的match进行过滤，distance越小越好，根据资料，能够小于最大distance的75%算好的匹配，参考资料：https://docs.opencv.org/4.x/dc/dc3/tutorial_py_matcher.html#autotoc_md1260
                List<DMatch> goodMatches = new ArrayList<>(matchList.size());
                double threshold = CollectionUtils.isEmpty(matchList) ? 999999 : 0.65 * Math.max(matchList.get(matchList.size() - 1).distance, 0.001); //完美匹配时distance会降到0，这里给一个很小的值保底让阈值不为0
                for (int i = 0; i < matchList.size(); ++i) {
                    DMatch m = matchList.get(i);
//                    if (m.distance < 65 && m.distance < threshold) { //小于某个值只是经验值，实际可能需要调整
                    if (m.distance < threshold) {
                        goodMatches.add(m);
                    } else {
                        break;
                    }
                }
                if (goodMatches.size() < 4) {
                    //不足4个点不能做图像变换
                    System.out.println("没有找到足够的匹配特征点");
                    return null;
                }
                int bestCnt = 100;
                DMatch[] bestMatches;
                if (goodMatches.size() > bestCnt) {
                    bestMatches = goodMatches.subList(0, bestCnt).toArray(new DMatch[0]);
                } else {
                    bestMatches = goodMatches.toArray(new DMatch[0]);
                }
                //计算最好的匹配点的平均距离（供研究）
                double distanceSum = 0;
                for (DMatch match : bestMatches) {
                    distanceSum += match.distance;
                }
                System.out.printf("最好的匹配点距离：%s\n", Arrays.stream(bestMatches).map(m -> String.valueOf(m.distance)).collect(Collectors.joining(",")));
                System.out.printf("最好的匹配点平均距离：%s\n", bestMatches.length > 0 ? String.valueOf(distanceSum / bestMatches.length) : "NaN");
                System.out.printf("最好的匹配点与特征点数最小值的比值：%f\n", (double) bestMatches.length / Math.min(kp1.rows(), kp2.rows()));
                Mat dst = new Mat();
//                MatOfDMatch bestMatOfMatches = new MatOfDMatch(bestMatches);
//                //绘制匹配特征点，drawMatches会将两个图片左右拼接起来，先用alpha=0的线让drawMatches的线透明
//                Features2d.drawMatches(img1, kp1, img2, kp2, bestMatOfMatches, dst, 1, new Scalar(255, 255, 255, 0), new Scalar(255, 255, 255), new MatOfByte(), Features2d.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS);
                drawMatches(dst, kp1, kp2, bestMatches, img1, img2);
//                drawMatches(dst, kp1, kp2, bestMatches, gray1, gray2);
                //调整图片大小以用于展示
                ImageUtils.fitImgSize(dst, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", dst, buffer);
                Platform.runLater(() -> {
                    imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray())));
//                    result.setText(String.format("发现匹配特征点%d个", goodMatches.size()));
                    result.setText(String.format("发现好的匹配特征点%d个", bestMatches.length));
                });
                return null;
            }
        };
        new Thread(task).start();
    }

    private void adaptImg(Mat img) {
        //如果不足3个通道则补成3通道的图像
        if (img.channels() < 3) {
            List<Mat> channels = new ArrayList<>(3);
            Core.split(img, channels);
            while (channels.size() < 3) {
                Mat newChannel = new Mat();
                channels.get(0).copyTo(newChannel);
                channels.add(newChannel);
            }
            Core.merge(channels, img);
        }
        //如果超过3个通道的情况下需要将4通道的图片改成3通道图片
        if (img.channels() > 3) {
            List<Mat> channels = new ArrayList<>();
            Core.split(img, channels);
            Core.merge(channels.subList(0, 3), img);
        }
    }

    private void drawMatches(Mat dst, MatOfKeyPoint kp1, MatOfKeyPoint kp2, DMatch[] ascSortedBestMatches, Mat queryImg, Mat trainImg) {
        KeyPoint[] keyPoints1 = kp1.toArray(), keyPoints2 = kp2.toArray();
//        //根据最好的几个特征点查找用于变换的矩阵
//        int sampleCnt = 4;
//        Point[] queryPoints = new Point[sampleCnt], trainPoints = new Point[sampleCnt];
//        for (int i = 0; i < sampleCnt; ++i) {
//            Point keyPt1 = keyPoints1[ascSortedBestMatches[i].queryIdx].pt;
//            Point keyPt2 = keyPoints2[ascSortedBestMatches[i].trainIdx].pt;
//            queryPoints[i] = new Point(Math.round(keyPt1.x), Math.round(keyPt1.y));
//            trainPoints[i] = new Point(Math.round(keyPt2.x), Math.round(keyPt2.y));
//        }
//        Mat M = Imgproc.getPerspectiveTransform(new MatOfPoint2f(queryPoints), new MatOfPoint2f(trainPoints));
////        Mat M = Calib3d.findHomography(new MatOfPoint2f(queryPoints), new MatOfPoint2f(trainPoints));
//        //进行图像变换
//        Imgproc.warpPerspective(queryImg, queryImg, M, new Size(queryImg.width(), queryImg.height()));
        //图片横向拼接，需要先将两个图片的高度对齐
        int diffH;
        if (queryImg.height() < trainImg.height()) {
            diffH = trainImg.height() - queryImg.height();
            Core.copyMakeBorder(queryImg, queryImg, 0, diffH, 0, 0, Core.BORDER_CONSTANT, new Scalar(0, 0, 0));
        } else if (queryImg.height() > trainImg.height()) {
            diffH = queryImg.height() - trainImg.height();
            Core.copyMakeBorder(trainImg, trainImg, 0, diffH, 0, 0, Core.BORDER_CONSTANT, new Scalar(0, 0, 0));
        }
        Imgproc.rectangle(queryImg, new Rect(new Point(0, 0), new Point(queryImg.width(), queryImg.height())), new Scalar(0, 0, 255), 5);
        Imgproc.rectangle(trainImg, new Rect(new Point(0, 0), new Point(trainImg.width(), trainImg.height())), new Scalar(255, 0, 0), 5);
        Core.hconcat(Arrays.asList(queryImg, trainImg), dst);
        //颜色池，从左往右是红色系到紫色系，建议匹配度越高的就用越接近红色的颜色
        List<Scalar> colorPool = Arrays.asList(new Scalar(49, 48, 214), new Scalar(85, 112, 225), new Scalar(110, 203, 253), new Scalar(227, 132, 9), new Scalar(231, 92, 108), new Scalar(54, 52, 45));
        //手动再对每个匹配的特征点用颜色池的颜色绘制连接线
        int i = 1;
        for (DMatch bestMatch : ascSortedBestMatches) {
//            Scalar color = colorPool.get((int) Math.floor(Math.random() * colorPool.size()));
            Scalar color = colorPool.get((int) Math.round(((double) i / ascSortedBestMatches.length) * 5));
            Point keyPoint1 = keyPoints1[bestMatch.queryIdx].pt;
            Point keyPoint2 = keyPoints2[bestMatch.trainIdx].pt;
            Point queryPt = new Point(keyPoint1.x, keyPoint1.y);
            Point trainPt = new Point(keyPoint2.x + queryImg.width(), keyPoint2.y); //keyPoint2在右侧，需要向右偏移img1图片的宽度
            Imgproc.circle(dst, queryPt, 8, color);
            Imgproc.circle(dst, trainPt, 8, color);
            Imgproc.line(dst, queryPt, trainPt, color, 2, Imgproc.LINE_AA);
//            Imgproc.putText(dst, String.valueOf(i), queryPt, Imgproc.FONT_HERSHEY_SIMPLEX, 2, new Scalar(0, 0, 255), 2, Imgproc.LINE_AA);
            i++;
        }
    }

    @Override
    public void start(Stage stage) {
        Text msg1 = new Text();
        Text msg2 = new Text();
        result = new Text();
        imgPath1 = "/Users/kwok/Downloads/protein_1.jpg";
        Button imageChooser1 = new Button("选择待比较图片");
        imageChooser1.setOnMouseClicked(e -> {
            File imgFile = chooseFile(stage);
            if (imgFile == null) return;
            imgPath1 = imgFile.getAbsolutePath();
            if (img1.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            msg1.setText("待比较图片已准备好");
        });
        imgPath2 = "/Users/kwok/Downloads/protein_2.jpg";
        Button imageChooser2 = new Button("选择基准图片");
        imageChooser2.setOnMouseClicked(e -> {
            File imgFile = chooseFile(stage);
            if (imgFile == null) return;
            imgPath2 = imgFile.getAbsolutePath();
            if (img2.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            msg2.setText("基准图片已准备好");
        });
        Button compareBtn = new Button("开始比较");
        compareBtn.setOnMouseClicked(event -> {
            if (StringUtils.isBlank(imgPath1) || StringUtils.isBlank(imgPath2)) {
                System.out.println("没有准备好2个图片，不能开始");
                return;
            }
            compareImages();
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser1, msg1, imageChooser2, msg2, compareBtn, result);
        bottomControl.setPrefHeight(50);
        imageView = new ImageView();
        BorderPane root = new BorderPane();
        root.setCenter(imageView);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("BF matcher example");
        stage.setResizable(false);
        stage.show();
    }

    private static File chooseFile(Stage stage) {
        FileChooser fileChooser = new FileChooser();
        fileChooser.setTitle("请选择图片");
        fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
        );
        fileChooser.getExtensionFilters().addAll(
                new FileChooser.ExtensionFilter("ALL", "*.*")
        );
        File pdfFile = fileChooser.showOpenDialog(stage);
        if (null == pdfFile) {
            System.out.println("用户取消选中文件");
            return null;
        }
        return pdfFile;
    }
}
